from python_ai.common.xcommon import *
import numpy as np

X = np.array([1., 2.]).reshape(1, -1)
state = [0., 0.]

w_state = np.array([[0.1, 0.2],
                    [0.3, 0.4]])
check_shape(w_state, 'w_state')
w_input = np.array([0.5, 0.6]).reshape(1, -1)
check_shape(w_input, 'w_input')
b = np.array([0.1, -0.1]).reshape(1, -1)
check_shape(b, 'b')

w_output = np.array([[1.], [2.]])
check_shape(w_output, 'w_output')
b_output = np.array([[0.1]])
check_shape(b_output, 'b_output')

for i in range(X.shape[1]):
    sep(f't-{i+1}')
    zt = np.dot(state, w_state) + np.dot(X[:, i:i+1], w_input) + b
    state = np.tanh(zt)
    output = np.dot(state, w_output) + b_output
    print(f'zt: {zt}')
    check_shape(zt, 'zt')
    print(f'state: {state}')
    check_shape(state, 'state')
    print(f'output: {output}')
    check_shape(output, 'output')
